package com.hehe.leo.hbase;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.client.coprocessor.AggregationClient;
import org.apache.hadoop.hbase.client.coprocessor.LongColumnInterpreter;
import org.apache.hadoop.hbase.coprocessor.AggregateImplementation;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.*;

/**
 * @author leo.jie (weixiao.me@aliyun.com)
 * @version 1.0
 * @organization CC
 * @website https://www.jlpyyf.com
 * @date 2019-10-10 21:44
 * @since 1.0
 */
public class HBaseRecordsCounter {
    private static final Map<String, String> ARGS_MAP = new HashMap<>();

    /**
     * 获取某张表的列簇
     *
     * @param conf      配置
     * @param tableName 表名
     * @return 列簇列表
     */
    private static List<String> getFamilies(Configuration conf, String tableName) {
        List<String> families = new ArrayList<>();
        Connection connection = null;
        try {
            connection = ConnectionFactory.createConnection(conf);
            Table table = connection.getTable(TableName.valueOf(tableName));
            Collection<HColumnDescriptor> familiesCollection = table.getTableDescriptor().getFamilies();
            if (null != familiesCollection && !familiesCollection.isEmpty()) {
                for (HColumnDescriptor family : familiesCollection) {
                    families.add(family.getNameAsString());
                }
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (null != connection) {
                try {
                    connection.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return families;
    }

    private static List<String> getAllTables(Configuration conf) {
        List<String> tables = new ArrayList<>();
        Connection connection = null;
        try {
            connection = ConnectionFactory.createConnection(conf);
            Admin hBaseAdmin = connection.getAdmin();
            for (HTableDescriptor hTableDescriptor : hBaseAdmin.listTables()) {
                if (hTableDescriptor.isMetaTable()) {
                    System.out.println("跳过表：" + hTableDescriptor.getNameAsString() + " 的统计！");
                    continue;
                }
                tables.add(hTableDescriptor.getNameAsString());
            }
            hBaseAdmin.close();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (null != connection) {
                try {
                    connection.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return tables;
    }

    /**
     * 为某张表添加协处理器
     *
     * @param conf      配置
     * @param tableName 表名
     */
    private static void addCoprocessor(Configuration conf, String tableName) {
        Connection connection = null;
        try {
            connection = ConnectionFactory.createConnection(conf);
            Admin hBaseAdmin = connection.getAdmin();
            Table table = connection.getTable(TableName.valueOf(tableName));
            TableName name = table.getName();
            HTableDescriptor htd = hBaseAdmin.getTableDescriptor(name);
            if (!htd.hasCoprocessor(AggregateImplementation.class.getName())) {
                hBaseAdmin.disableTable(name);
                htd.addCoprocessor(AggregateImplementation.class.getName());
                hBaseAdmin.modifyTable(name, htd);
                hBaseAdmin.enableTable(name);
                System.out.println("table:" + tableName + " 的协处理器添加成功！");
            } else {
                System.out.println("table:" + tableName + " 的协处理器已经被添加！");
            }
            hBaseAdmin.close();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (null != connection) {
                try {
                    connection.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 执行统计行数的逻辑
     *
     * @param conf      配置
     * @param tableName 表名
     * @param family    列簇
     */
    private static void execCount(Configuration conf, String tableName, String family) {
        AggregationClient aggregationClient = new AggregationClient(conf);
        Scan scan = new Scan();
        scan.addFamily(Bytes.toBytes(family));
        long start = System.currentTimeMillis();
        try {
            long rowCount = aggregationClient.rowCount(TableName.valueOf(tableName),
                    new LongColumnInterpreter(), scan);
            System.out.println("Count Table:" + tableName + " Family:" + family
                    + " Row Count Is: " + rowCount + " Time Cost: " +
                    (System.currentTimeMillis() - start) + " ms");
        } catch (Throwable throwable) {
            throwable.printStackTrace();
        }
    }

    /**
     * 转换main函数参数
     *
     * @param args main函数参数
     */
    private static void parseArgs(String[] args) {
        if (args == null || args.length == 0) {
            throw new IndexOutOfBoundsException("参数不能为空！");
        }
        for (String arg : args) {
            if (arg.contains("=")) {
                String[] tmpArgs = arg.split("=");
                if (tmpArgs.length != 2) {
                    continue;
                }
                if (!tmpArgs[0].startsWith("-")) {
                    continue;
                }
                if (tmpArgs[1] == null || tmpArgs[1].trim().length() == 0) {
                    continue;
                }
                ARGS_MAP.put(tmpArgs[0].substring(1), tmpArgs[1]);
            }
        }
    }

    private static String getStrValue(String key) {
        if (!ARGS_MAP.containsKey(key)) {
            printUsage();
            throw new RuntimeException("不存在的配置：" + key);
        }
        return ARGS_MAP.get(key);
    }

    private static String getStrValueWithDefault(String key, String defaultValue) {
        return ARGS_MAP.getOrDefault(key, defaultValue);
    }

    private static long getLongValueWithDefault(String key, long defaultValue) {
        return Long.parseLong(getStrValueWithDefault(key, String.valueOf(defaultValue)));
    }

    private static void printUsage() {
        System.out.println("-table=tableName\t必须参数\t需要统计的表名;\n" +
                "-zk=zkAddress\t必须参数\t需要的zookeeper地址;\n" +
                "-rpc.timeout=6000\t默认6000ms\t指定rpc通信的超时时间;\n" +
                "-scanner=1000\t默认1000\t单次scanner操作缓存的行数;\n" +
                "-exclude=t1,t2\t非必须参数，不选则统计所有表\t统计需要排除的表名;");
    }

    public static void main(String[] args) {
        parseArgs(args);
        String zk = getStrValue("zk");
        String tableName = getStrValueWithDefault("table", null);
        long rpcTimeOut = getLongValueWithDefault("rpc.timeout", 6000);
        long scanCache = getLongValueWithDefault("scanner", 1000);
        String excludeTable = getStrValueWithDefault("exclude", null);
        List<String> excludeTables = null;
        if (null != excludeTable) {
            excludeTables = new ArrayList<>(Arrays.asList(excludeTable.split(",")));
        }
        Configuration conf = new Configuration();
        conf.set("hbase.zookeeper.quorum", zk);
        conf.setLong("hbase.rpc.timeout", rpcTimeOut);
        // 设置Scan缓存
        conf.setLong("hbase.client.scanner.caching", scanCache);
        List<String> allTables = new ArrayList<>();
        if (null != tableName) {
            allTables.add(tableName);
        } else {
            allTables = getAllTables(conf);
        }
        if (null != excludeTables && !excludeTables.isEmpty()) {
            allTables.removeAll(excludeTables);
        }
        for (String table : allTables) {
            addCoprocessor(conf, table);
            List<String> families = getFamilies(conf, table);
            if (null == families || families.isEmpty()) {
                continue;
            }
            for (String family : families) {
                execCount(conf, table, family);
            }
        }

    }
}